#!/usr/bin/env python
# -*- coding:utf-8 -*-
from langchain_text_splitters import CharacterTextSplitter, MarkdownTextSplitter, RecursiveCharacterTextSplitter


class DocumentSplitter:
    chunk_size = 2048
    chunk_overlap = 10
    separator = ["\n\n"]
    add_start_index = True

    @classmethod
    def load_txt_splitter(cls, docs):
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=cls.chunk_size, chunk_overlap=cls.chunk_overlap, add_start_index=cls.add_start_index)
        split_docs = text_splitter.split_documents(docs)

        return split_docs

    @classmethod
    def load_md_splitter(cls, docs):
        text_splitter = MarkdownTextSplitter(chunk_size=cls.chunk_size, chunk_overlap=cls.chunk_overlap, add_start_index=cls.add_start_index)
        split_docs = text_splitter.split_documents(docs)

        return split_docs

    @classmethod
    def load_pdf_splitter(cls, docs):
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=cls.chunk_size, chunk_overlap=cls.chunk_overlap, add_start_index=cls.add_start_index)
        split_docs = text_splitter.split_documents(docs)

        return split_docs

    @classmethod
    def load_jpg_splitter(cls, docs):
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=cls.chunk_size, chunk_overlap=cls.chunk_overlap, add_start_index=cls.add_start_index)
        split_docs = text_splitter.create_documents([docs])

        return split_docs

    @classmethod
    def load_docx_splitter(cls, docs):
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=cls.chunk_size, chunk_overlap=cls.chunk_overlap, add_start_index=cls.add_start_index)
        split_docs = text_splitter.split_documents(docs)

        return split_docs

    @classmethod
    def load_ppt_splitter(cls, docs):
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=cls.chunk_size, chunk_overlap=cls.chunk_overlap, add_start_index=cls.add_start_index)
        split_docs = text_splitter.split_documents(docs)

        return split_docs

    @classmethod
    def load_csv_splitter(cls, docs):
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=cls.chunk_size, chunk_overlap=cls.chunk_overlap, add_start_index=cls.add_start_index)
        split_docs = text_splitter.split_documents(docs)

        return split_docs